import os
import glob
import tensorflow as tf
from algorithms.Deep3DFaceReconstruction.load_data import *
import dlib


class Deep3DFaceRecon:
    def __init__(self):
        # input and output folder
        self.save_path = 'algorithms/Deep3DFaceReconstruction/output'

        image_path = 'input'
        self.img_list = glob.glob(
            'algorithms/Deep3DFaceReconstruction/' + image_path + '/' + '*.png')

        # read BFM face model
        # transfer original BFM model to our model
        if not os.path.isfile('algorithms/Deep3DFaceReconstruction/BFM/BFM_model_front.mat'):
            transferBFM09()

        # read face model
        self.facemodel = BFM()
        # read standard landmarks for preprocessing images
        self.lm3D = load_lm3d()

        self.graph_def = self.load_graph(
            'algorithms/Deep3DFaceReconstruction/network/FaceReconModel.pb')

        self.dlib_detector = dlib.get_frontal_face_detector()
        self.dlib_predictor = dlib.shape_predictor(
            'models/dlib_featurepoints/predictor_68_face_landmarks.dat')

    def load_graph(self, graph_filename):
        with tf.gfile.GFile(graph_filename, 'rb') as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())

        return graph_def
